#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
QTorch量化交易框架 - 因子扩展演示
展示如何使用基本面因子和另类数据因子
"""

import logging
import argparse
import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from datetime import datetime, timedelta

from qtorch.data.aksharestockerfetcher import AKShareStockFetcher
from qtorch.factor.qfactor import QFactor, Factor, MultiFactorModel
from qtorch.factor.fundamental_factors import FundamentalFactors, AKShareFinancialDataProvider, TushareFinancialDataProvider
from qtorch.factor.alternative_factors import AlternativeFactors, TextAnalyzer
from qtorch.core.qengine import QEngine

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler()]
)

def load_data(symbol='000001', start_date='2020-01-01', end_date='2022-12-31'):
    """加载市场数据"""
    data_fetcher = AKShareStockFetcher(cache_dir='./data_cache')
    data = data_fetcher.get_market_data(symbol, start_date, end_date)
    
    if data is None or data.empty:
        raise ValueError(f"获取 {symbol} 数据失败")
        
    logging.info(f"成功加载 {symbol} 数据，共 {len(data)} 个交易日")
    return data

def fundamental_factor_demo(symbol='000001', start_date='2020-01-01', end_date='2022-12-31'):
    """基本面因子演示"""
    logging.info("开始基本面因子演示...")
    
    # 创建输出目录
    os.makedirs('./output/fundamental', exist_ok=True)
    
    # 加载价格数据
    price_data = load_data(symbol, start_date, end_date)
    
    # 初始化QFactor和基本面因子库
    qfactor = QFactor()
    fundamental = FundamentalFactors(data_provider="akshare")
    
    try:
        # 创建基本面因子
        logging.info("计算PE因子...")
        pe_factor = fundamental.create_pe_factor(symbol, price_data)
        
        logging.info("计算PB因子...")
        pb_factor = fundamental.create_pb_factor(symbol, price_data)
        
        logging.info("计算ROE因子...")
        roe_factor = fundamental.create_roe_factor(symbol, price_data)
        
        logging.info("计算股息率因子...")
        dividend_yield_factor = fundamental.create_dividend_yield_factor(symbol, price_data)
        
        # 创建自定义因子并添加到因子库
        if pe_factor is not None:
            pe_factor_obj = Factor(name="pe_ttm", data=pe_factor, description="市盈率(TTM)", category="估值")
            qfactor.library.add_factor(pe_factor_obj)
        
        if pb_factor is not None:
            pb_factor_obj = Factor(name="pb", data=pb_factor, description="市净率", category="估值")
            qfactor.library.add_factor(pb_factor_obj)
        
        if roe_factor is not None:
            roe_factor_obj = Factor(name="roe", data=roe_factor, description="净资产收益率", category="质量")
            qfactor.library.add_factor(roe_factor_obj)
            
        if dividend_yield_factor is not None:
            dividend_factor_obj = Factor(name="dividend_yield", data=dividend_yield_factor, description="股息率", category="股息")
            qfactor.library.add_factor(dividend_factor_obj)
            
        # 计算未来收益率作为目标变量
        future_returns = price_data['close'].pct_change(20).shift(-20)  # 20日后收益率
        
        # 因子分析
        for factor_name in ["pe_ttm", "pb", "roe", "dividend_yield"]:
            factor = qfactor.get_factor(factor_name)
            if factor is not None and factor.data is not None:
                # 计算因子IC和回测
                logging.info(f"分析因子: {factor_name}")
                ic_result = factor.calc_ic(future_returns)
                backtest_result = factor.backtest(future_returns)
                
                # 绘制因子分析图表
                fig = factor.plot(future_returns)
                plt.savefig(f'./output/fundamental/{factor_name}_analysis.png')
                plt.close(fig)
                
                # 打印因子IC值
                for k, v in ic_result.items():
                    if isinstance(v, float):
                        logging.info(f"{factor_name} {k}: {v:.4f}")
        
        # 因子相关性分析
        logging.info("分析因子相关性...")
        corr_matrix = qfactor.correlation_matrix()
        
        if not corr_matrix.empty:
            plt.figure(figsize=(10, 8))
            sns.heatmap(corr_matrix, annot=True, cmap='coolwarm', vmin=-1, vmax=1, linewidths=0.5)
            plt.title("基本面因子相关性矩阵")
            plt.tight_layout()
            plt.savefig('./output/fundamental/factor_correlation.png')
            plt.close()
        
        # 多因子模型
        logging.info("构建多因子模型...")
        
        # 创建因子模型
        model = qfactor.create_model("基本面多因子模型")
        
        # 添加因子
        for factor_name in ["pe_ttm", "pb", "roe", "dividend_yield"]:
            factor = qfactor.get_factor(factor_name)
            if factor is not None and factor.data is not None:
                model.add_factor(factor)
        
        if len(model.factors) > 0:
            # 使用等权重组合因子
            combined_factor = model.combine_factors(method='equal_weight')
            
            # 回测组合因子
            backtest_result = model.backtest(future_returns)
            
            # 绘制多因子模型分析图表
            fig = model.plot(future_returns)
            plt.savefig('./output/fundamental/multi_factor_model.png')
            plt.close(fig)
            
            logging.info(f"多因子模型年化收益率: {model.factors.get('annual_long_short', 0):.4f}")
            
        return {
            'qfactor': qfactor,
            'fundamental': fundamental,
            'price_data': price_data,
            'factors': {
                'pe_ttm': pe_factor,
                'pb': pb_factor,
                'roe': roe_factor,
                'dividend_yield': dividend_yield_factor
            },
            'model': model if 'model' in locals() else None
        }
    
    except Exception as e:
        logging.error(f"基本面因子演示失败: {str(e)}", exc_info=True)
        return None

def alternative_factor_demo(symbol='000001', keyword=None, start_date='2020-01-01', end_date='2022-12-31'):
    """另类数据因子演示"""
    logging.info("开始另类数据因子演示...")
    
    # 创建输出目录
    os.makedirs('./output/alternative', exist_ok=True)
    
    # 加载价格数据
    price_data = load_data(symbol, start_date, end_date)
    
    # 如果没有提供关键词，使用股票代码
    if keyword is None:
        keyword = f"股票{symbol}"
        
    # 初始化QFactor和另类数据因子库
    qfactor = QFactor()
    alternative = AlternativeFactors()
    
    try:
        # 创建另类数据因子
        logging.info("计算搜索指数因子...")
        search_factor = alternative.create_search_index_factor(keyword, start_date, end_date, price_data)
        
        logging.info("计算新闻情感因子...")
        news_factor = alternative.create_news_sentiment_factor(symbol, start_date, end_date, price_data)
        
        logging.info("计算社交情感因子...")
        social_factor = alternative.create_social_sentiment_factor(keyword, start_date, end_date, price_data)
        
        logging.info("计算公告因子...")
        announcement_factor = alternative.create_announcement_factor(symbol, start_date, end_date, price_data)
        
        # 创建自定义因子并添加到因子库
        if search_factor is not None:
            search_factor_obj = Factor(name="search_index", data=search_factor, description="搜索指数", category="关注度")
            qfactor.library.add_factor(search_factor_obj)
        
        if news_factor is not None:
            news_factor_obj = Factor(name="news_sentiment", data=news_factor, description="新闻情感", category="情感")
            qfactor.library.add_factor(news_factor_obj)
        
        if social_factor is not None:
            social_factor_obj = Factor(name="social_sentiment", data=social_factor, description="社交情感", category="情感")
            qfactor.library.add_factor(social_factor_obj)
            
        if announcement_factor is not None:
            announcement_factor_obj = Factor(name="announcement_factor", data=announcement_factor, description="公告因子", category="事件")
            qfactor.library.add_factor(announcement_factor_obj)
            
        # 计算未来收益率作为目标变量
        future_returns = price_data['close'].pct_change(5).shift(-5)  # 5日后收益率
        
        # 因子分析
        for factor_name in ["search_index", "news_sentiment", "social_sentiment", "announcement_factor"]:
            factor = qfactor.get_factor(factor_name)
            if factor is not None and factor.data is not None:
                # 计算因子IC和回测
                logging.info(f"分析因子: {factor_name}")
                ic_result = factor.calc_ic(future_returns)
                backtest_result = factor.backtest(future_returns)
                
                # 绘制因子分析图表
                fig = factor.plot(future_returns)
                plt.savefig(f'./output/alternative/{factor_name}_analysis.png')
                plt.close(fig)
                
                # 打印因子IC值
                for k, v in ic_result.items():
                    if isinstance(v, float):
                        logging.info(f"{factor_name} {k}: {v:.4f}")
        
        # 因子相关性分析
        logging.info("分析因子相关性...")
        corr_matrix = qfactor.correlation_matrix()
        
        if not corr_matrix.empty:
            plt.figure(figsize=(10, 8))
            sns.heatmap(corr_matrix, annot=True, cmap='coolwarm', vmin=-1, vmax=1, linewidths=0.5)
            plt.title("另类数据因子相关性矩阵")
            plt.tight_layout()
            plt.savefig('./output/alternative/factor_correlation.png')
            plt.close()
        
        # 多因子模型
        logging.info("构建另类数据多因子模型...")
        
        # 创建因子模型
        model = qfactor.create_model("另类数据多因子模型")
        
        # 添加因子
        for factor_name in ["search_index", "news_sentiment", "social_sentiment", "announcement_factor"]:
            factor = qfactor.get_factor(factor_name)
            if factor is not None and factor.data is not None:
                model.add_factor(factor)
        
        if len(model.factors) > 0:
            # 使用等权重组合因子
            combined_factor = model.combine_factors(method='equal_weight')
            
            # 回测组合因子
            backtest_result = model.backtest(future_returns)
            
            # 绘制多因子模型分析图表
            fig = model.plot(future_returns)
            plt.savefig('./output/alternative/multi_factor_model.png')
            plt.close(fig)
            
            logging.info(f"多因子模型年化收益率: {model.factors.get('annual_long_short', 0):.4f}")
        
        # 文本分析演示
        logging.info("演示文本分析功能...")
        
        # 示例文本
        positive_text = "公司业绩持续增长，营收和利润均超出市场预期，未来发展前景良好。"
        negative_text = "公司业绩不及预期，营收下滑，面临较大的经营压力和竞争挑战。"
        neutral_text = "公司发布季度财报，营收同比持平，符合市场预期。"
        
        # 情感分析
        positive_score = alternative.analyze_text_sentiment(positive_text)
        negative_score = alternative.analyze_text_sentiment(negative_text)
        neutral_score = alternative.analyze_text_sentiment(neutral_text)
        
        logging.info(f"积极文本情感得分: {positive_score:.4f}")
        logging.info(f"消极文本情感得分: {negative_score:.4f}")
        logging.info(f"中性文本情感得分: {neutral_score:.4f}")
        
        # 关键词提取
        keywords = alternative.extract_text_keywords(positive_text + " " + negative_text, top_n=10)
        logging.info(f"提取的关键词: {keywords}")
            
        return {
            'qfactor': qfactor,
            'alternative': alternative,
            'price_data': price_data,
            'factors': {
                'search_index': search_factor,
                'news_sentiment': news_factor,
                'social_sentiment': social_factor,
                'announcement_factor': announcement_factor
            },
            'model': model if 'model' in locals() else None
        }
    
    except Exception as e:
        logging.error(f"另类数据因子演示失败: {str(e)}", exc_info=True)
        return None

def combined_factor_model_demo(symbol='000001', start_date='2020-01-01', end_date='2022-12-31'):
    """组合因子模型演示"""
    logging.info("开始组合因子模型演示...")
    
    # 创建输出目录
    os.makedirs('./output/combined', exist_ok=True)
    
    try:
        # 运行基本面因子演示
        fund_result = fundamental_factor_demo(symbol, start_date, end_date)
        
        # 运行另类数据因子演示
        alt_result = alternative_factor_demo(symbol, None, start_date, end_date)
        
        if fund_result is None or alt_result is None:
            logging.error("无法构建组合模型，因子演示失败")
            return None
            
        # 初始化QFactor
        qfactor = QFactor()
        
        # 收集所有因子
        all_factors = {}
        
        # 添加基本面因子
        for name, factor in fund_result['factors'].items():
            if factor is not None:
                factor_obj = Factor(name=name, data=factor, description=name, category="基本面")
                qfactor.library.add_factor(factor_obj)
                all_factors[name] = factor
                
        # 添加另类数据因子
        for name, factor in alt_result['factors'].items():
            if factor is not None:
                factor_obj = Factor(name=name, data=factor, description=name, category="另类数据")
                qfactor.library.add_factor(factor_obj)
                all_factors[name] = factor
        
        # 计算未来收益率作为目标变量
        price_data = fund_result['price_data']
        future_returns = price_data['close'].pct_change(10).shift(-10)  # 10日后收益率
        
        # 因子相关性分析
        logging.info("分析因子相关性...")
        corr_matrix = qfactor.correlation_matrix()
        
        if not corr_matrix.empty:
            plt.figure(figsize=(12, 10))
            sns.heatmap(corr_matrix, annot=True, cmap='coolwarm', vmin=-1, vmax=1, linewidths=0.5)
            plt.title("组合因子相关性矩阵")
            plt.tight_layout()
            plt.savefig('./output/combined/factor_correlation.png')
            plt.close()
        
        # 创建组合因子模型
        logging.info("构建组合因子模型...")
        model = qfactor.create_model("组合因子模型")
        
        # 添加因子
        for name, factor in all_factors.items():
            if factor is not None:
                factor_obj = qfactor.get_factor(name)
                if factor_obj is not None:
                    model.add_factor(factor_obj)
        
        if len(model.factors) > 0:
            # 使用等权重组合因子
            combined_factor = model.combine_factors(method='equal_weight')
            
            # 回测组合因子
            backtest_result = model.backtest(future_returns)
            
            # 绘制多因子模型分析图表
            fig = model.plot(future_returns)
            plt.savefig('./output/combined/combined_factor_model.png')
            plt.close(fig)
            
            # 使用回归方法组合因子
            try:
                model.fit_regression(future_returns, model_type='ridge', alpha=1.0)
                
                # 回测回归组合因子
                regression_result = model.backtest(future_returns)
                
                # 绘制回归模型分析图表
                fig = model.plot(future_returns)
                plt.savefig('./output/combined/regression_factor_model.png')
                plt.close(fig)
                
                # 打印因子权重
                logging.info("回归模型因子权重:")
                for name, weight in model.weights.items():
                    logging.info(f"{name}: {weight:.4f}")
            except Exception as e:
                logging.error(f"回归模型构建失败: {str(e)}")
        
        return {
            'qfactor': qfactor,
            'all_factors': all_factors,
            'model': model,
            'price_data': price_data
        }
    
    except Exception as e:
        logging.error(f"组合因子模型演示失败: {str(e)}", exc_info=True)
        return None

def factor_backtest_demo(symbol='000001', start_date='2020-01-01', end_date='2022-12-31'):
    """因子选股回测演示"""
    logging.info("开始因子选股回测演示...")
    
    # 创建输出目录
    os.makedirs('./output/backtest', exist_ok=True)
    
    try:
        # 运行组合因子模型演示
        combined_result = combined_factor_model_demo(symbol, start_date, end_date)
        
        if combined_result is None:
            logging.error("无法执行因子选股回测，组合因子模型演示失败")
            return None
        
        # 提取模型和价格数据
        model = combined_result['model']
        price_data = combined_result['price_data']
        
        if model is None or price_data is None:
            logging.error("模型或价格数据为空")
            return None
            
        # 使用QEngine进行回测
        logging.info("使用QEngine进行因子选股回测...")
        
        # 创建和定义基于因子的策略
        from qtorch.core.qstrategy import QStrategy
        
        class FactorStrategy(QStrategy):
            """基于因子的选股策略"""
            
            def __init__(self, factor_data, rebalance_freq=20):
                super().__init__()
                self.factor_data = factor_data
                self.factor_values = None
                self.rebalance_freq = rebalance_freq
                self.last_rebalance_day = 0
                self.day_count = 0
                
            def next(self):
                # 更新天数计数
                self.day_count += 1
                
                # 检查是否是当前数据点的索引
                current_date = self.data.name
                
                # 仅在再平衡日进行交易
                if (self.day_count - self.last_rebalance_day) < self.rebalance_freq:
                    return 0  # 持有当前仓位
                
                # 更新再平衡日
                self.last_rebalance_day = self.day_count
                
                # 获取当前日期的因子值
                if current_date in self.factor_data.index:
                    factor_value = self.factor_data.loc[current_date]
                    
                    # 根据因子值生成交易信号
                    if factor_value > 0.5:  # 因子值较高，看多
                        return 1  # 买入信号
                    elif factor_value < -0.5:  # 因子值较低，看空
                        return -1  # 卖出信号
                
                return 0  # 默认持有
        
        # 获取组合因子值
        factor_data = model.combined_factor
        
        # 确保因子数据和价格数据的索引类型一致
        if not isinstance(factor_data.index, pd.DatetimeIndex):
            factor_data.index = pd.to_datetime(factor_data.index)
        
        if not isinstance(price_data.index, pd.DatetimeIndex):
            price_data.index = pd.to_datetime(price_data.index)
        
        # 创建引擎
        engine = QEngine()
        engine.add_data(symbol, start_date, end_date)
        
        # 添加因子策略
        engine.add_strategy(FactorStrategy, factor_data=factor_data, rebalance_freq=20)
        
        # 运行回测
        result = engine.run()
        
        # 绘制回测结果
        result.plot(output_dir='./output/backtest')
        
        # 打印回测指标
        logging.info("因子选股回测结果:")
        for key, value in result.metrics.items():
            if key in ["总收益率", "年化收益率", "最大回撤", "夏普比率", "索提诺比率", "胜率", "盈亏比"]:
                if isinstance(value, float):
                    logging.info(f"{key}: {value:.4f}")
                else:
                    logging.info(f"{key}: {value}")
        
        return {
            'model': model,
            'factor_data': factor_data,
            'price_data': price_data,
            'backtest_result': result
        }
    
    except Exception as e:
        logging.error(f"因子选股回测演示失败: {str(e)}", exc_info=True)
        return None

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='QTorch量化交易框架 - 因子扩展演示')
    parser.add_argument('--mode', type=str, default='all', 
                        choices=['fundamental', 'alternative', 'combined', 'backtest', 'all'],
                        help='演示模式')
    parser.add_argument('--symbol', type=str, default='000001',
                        help='股票代码')
    parser.add_argument('--keyword', type=str, default=None,
                        help='搜索关键词，如不提供则使用股票代码')
    parser.add_argument('--start', type=str, default='2020-01-01',
                        help='回测开始日期')
    parser.add_argument('--end', type=str, default='2022-12-31',
                        help='回测结束日期')
                        
    args = parser.parse_args()
    
    # 创建输出目录
    os.makedirs('./output', exist_ok=True)
    
    try:
        if args.mode == 'fundamental' or args.mode == 'all':
            fundamental_factor_demo(args.symbol, args.start, args.end)
            
        if args.mode == 'alternative' or args.mode == 'all':
            alternative_factor_demo(args.symbol, args.keyword, args.start, args.end)
            
        if args.mode == 'combined' or args.mode == 'all':
            combined_factor_model_demo(args.symbol, args.start, args.end)
            
        if args.mode == 'backtest' or args.mode == 'all':
            factor_backtest_demo(args.symbol, args.start, args.end)
            
        logging.info(f"所有演示任务已完成！结果保存在: ./output")
            
    except Exception as e:
        logging.error(f"演示过程中出错: {str(e)}", exc_info=True)
        raise

if __name__ == '__main__':
    main()